import sys,os,time,subprocess
from collections import defaultdict
import xlrd,glob
from datetime import datetime

d = time.strftime("%Y%m%d",time.localtime())
site = {'changsha':'CS','guangzhou':'GZ','zhengzhou':'ZZ','beijing':'BJ','fuzhou':'FZ','shanghai':'SH'}

def dealkk(indir,sampleid,kk,the_list,i,autores,sampledata,sampletype):
    sss = defaultdict(int)
    with open(f'{indir}/list/{sampleid[i+4].split(".")[0]}原始数据列表.kk.xls','w',encoding = 'gbk') as out:
        out.write(f'\t\t\t\t{sampletype[i+4]}\n')
        out.write(f'\t\t\t\t{sampledata[i+4]}\n')
        out.write(f'D\tG\tS\tspecies_zn\t{sampleid[i+4].split(".")[0].split("-")[0]}\n')
        for line in range(3, kk.nrows):
            l = [str(j.value).strip() for j in kk.row(line)]
            species = ' '.join(l[2].split('_')[2:])
            if species in the_list or l[2] in autores[sampleid[i+4]]:
                if int(float(l[i+4])) > 0:
                    s = f'{l[0]}\t{l[1]}\t{l[2]}\t{l[3]}'
                    sss[s] = int(float(l[i+4]))
        allssslistsorted = sorted(sss.items(), key = lambda x: x[1], reverse = True)
        for anytuple in allssslistsorted:
            out.write(f'{anytuple[0]}\t{str(anytuple[1])}\n')
    out.close

def dealnt(indir,sampleid,kk,the_list,i,autores,sampledata,sampletype):
    sss = defaultdict(int)
    with open(f'{indir}/list/{sampleid[i+4].split(".")[0]}原始数据列表.nt.xls','w',encoding = 'gbk') as out:
        out.write(f'\t\t\t\t{sampletype[i+4]}\n')
        out.write(f'\t\t\t\t{sampledata[i+4]}\n')
        out.write(f'D\tG\tS\tspecies_zn\t{sampleid[i+4].split(".")[0].split("-")[0]}\n')
        for line in range(3, kk.nrows):
            l = [str(j.value).strip() for j in kk.row(line)]
            if l[2].find('|') > -1:
                species = l[2].split('|')[1]
            else:
                continue
            if species in the_list or l[2] in autores[sampleid[i+4].split('.')[0]]:
                if int(float(l[i+4])) > 0:
                    s = f'{l[0]}\t{l[1]}\t{l[2]}\t{l[3]}'
                    sss[s] = int(float(l[i+4]))
        allssslistsorted = sorted(sss.items(), key = lambda x: x[1], reverse = True)
        for anytuple in allssslistsorted:
            out.write(f'{anytuple[0]}\t{str(anytuple[1])}\n')
    out.close

def getlist(indir, _sampletype, _Microbe):
    s = indir.split('/')[-2]
    kkfile = glob.glob(f'{indir}/*.mergeKK.xlsx')[0]
    ntfile = glob.glob(f'{indir}/*.mergeNT.xlsx')[0]
    autofiles = glob.glob(f'/data/pub/today/.auto/{s}/{sys.argv[2]}/output/*.xlsx')
    autofile = glob.glob(f'{indir}/kkbwa/*.autoTest.*.xlsx')[0]
    sa_type = defaultdict(str)
    Balf,Blood,CSF,FS,NY,Other,XS = [],[],[],[],[],[],[]

    with open(_sampletype) as sample_type:
        sample_type.readline()
        for l in sample_type:
            i = l.strip().split('\t')
            sa_type[i[0]] = i[1]
    
    book = xlrd.open_workbook(_Microbe)
    listsheet = book.sheet_by_index(0)
    for i in range(1, listsheet.nrows):
        l = [str(j.value).strip() for j in listsheet.row(i)]
        if l[13].find('背景') > -1:
            Balf.append(l[8])
        if l[14].find('背景') > -1:
            XS.append(l[8])
        if l[15].find('背景') > -1:
            FS.append(l[8])
        if l[16].find('背景') > -1:
            NY.append(l[8])
        if l[17].find('背景') > -1:
            Other.append(l[8])
        if l[18].find('背景') > -1:
            CSF.append(l[8])
        if l[20].find('背景') > -1:
            Blood.append(l[8])
    while 1:
        print(kkfile,ntfile,autofiles)
        if os.path.exists(kkfile) and os.path.exists(ntfile) and len(autofiles) == 2:
            break
        time.sleep(60)

    autores = defaultdict(list)
    auto = xlrd.open_workbook(autofile)
    autosheet = auto.sheet_by_index(0)
    for i in range(1, autosheet.nrows):
        l = [str(j.value).strip() for j in autosheet.row(i)]
        if l[9] == '高':
            autores[l[3]].append(l[6])

    kf = xlrd.open_workbook(kkfile)
    kksheet = kf.sheet_by_index(0)
    sampletype = [str(j.value).strip() for j in kksheet.row(0)]
    sampledata = [str(j.value).strip() for j in kksheet.row(1)]
    sampleid = [str(j.value).strip() for j in kksheet.row(2)]
    for i,e in enumerate(sampletype[4:-1]):
        if sa_type[e] == 'BALF':
            dealkk(indir,sampleid,kksheet,Balf,i,autores,sampledata,sampletype)
        if sa_type[e] == 'Blood':
            dealkk(indir,sampleid,kksheet,Blood,i,autores,sampledata,sampletype)
        if sa_type[e] == 'XS':
            dealkk(indir,sampleid,kksheet,XS,i,autores,sampledata,sampletype)
        if sa_type[e] == 'FS':
            dealkk(indir,sampleid,kksheet,FS,i,autores,sampledata,sampletype)
        if sa_type[e] == 'NY':
            dealkk(indir,sampleid,kksheet,NY,i,autores,sampledata,sampletype)
        if sa_type[e] == 'CSF':
            dealkk(indir,sampleid,kksheet,CSF,i,autores,sampledata,sampletype)
        if sa_type[e] == 'Other':
            dealkk(indir,sampleid,kksheet,Other,i,autores,sampledata,sampletype)

    nf = xlrd.open_workbook(ntfile)
    ntsheet = nf.sheet_by_index(0)
    sampletype = [str(j.value).strip() for j in ntsheet.row(0)]
    sampleid = [str(j.value).strip() for j in ntsheet.row(2)]
    for i,e in enumerate(sampletype[4:-1]):
        if sa_type[e] == 'BALF':
            dealnt(indir,sampleid,ntsheet,Balf,i,autores,sampledata,sampletype)
        if sa_type[e] == 'Blood':
            dealnt(indir,sampleid,ntsheet,Blood,i,autores,sampledata,sampletype)
        if sa_type[e] == 'XS':
            dealnt(indir,sampleid,ntsheet,XS,i,autores,sampledata,sampletype)
        if sa_type[e] == 'FS':
            dealnt(indir,sampleid,ntsheet,FS,i,autores,sampledata,sampletype)
        if sa_type[e] == 'NY':
            dealnt(indir,sampleid,ntsheet,NY,i,autores,sampledata,sampletype)
        if sa_type[e] == 'CSF':
            dealnt(indir,sampleid,ntsheet,CSF,i,autores,sampledata,sampletype)
        if sa_type[e] == 'Other':
            dealnt(indir,sampleid,ntsheet,Other,i,autores,sampledata,sampletype)

if __name__ == '__main__':
    inputdir = f'{sys.argv[1]}/{sys.argv[2]}'
    the_type = f'/data/softwares/mngs_scripts/auto/autoPathoSeqPy/tag/PathoSeq_SampleType.txt'
    microbe = f'/data/softwares/mngs_scripts/auto/autoPathoSeqPy/tag/PathoSeq_Microbe.Py_20210721.xlsx'
    runid = sys.argv[2]
    getlist(inputdir,the_type,microbe)
    script = f'''cd {sys.argv[1]} && tar zcvf {runid}.tar.gz {runid}'''
    subprocess.run(script,shell = True)
